﻿using CrazyMVSamLibrary.Base;
using Emgu.CV.Structure;
using Emgu.CV;
using System.Collections.Generic;
using System.Drawing;
using CrazyMVSamLibrary.Model;

namespace CrazyMVSamLibrary.Function
{
    public class SamFunction
    {
        private Sam sam;
        private Transforms ts;
        private CrazyImageEncode image;

        public SamFunction()
        {
            this.sam = new Sam();
            this.ts = new Transforms(1024);
            this.image = null;
        }

        public void LoadModel(string encodeModelPath, string decodeModelPath)
        {
            this.sam.LoadONNXModel(encodeModelPath, decodeModelPath);
        }

        public void EncodeImage(byte[,,] data, int width, int height)
        {
            Mat img;
            if (data.GetLongLength(data.Rank - 1) == 1)
            {
                img = new Image<Gray, byte>(data).Mat;
            }
            else
            {
                img = new Image<Bgr, byte>(data).Mat;
            }
            float[] mImgEmbedding = this.sam.Encode(this.ts, img, width, height);
            this.image = new CrazyImageEncode
            {
                Data = data,
                ImageEmbedding = mImgEmbedding,
                Width = width,
                Height = height
            };
        }

        public CrazyResult PointsDecode(List<Point> addPoints, List<Point> removePoints)
        {
            int addLength = addPoints.Count;
            int removeLength = removePoints.Count;
            int width = this.image.Width;
            int height = this.image.Height;
            List<Promotion> mPromotionList = new List<Promotion>();
            for (int i = 0; i < addLength; i++)
            {
                PointPromotion point = new PointPromotion(OperatorType.ADD)
                {
                    X = addPoints[i].X,
                    Y = addPoints[i].Y
                };
                PointPromotion ptn = ts.ApplyCoords(point, width, height);
                mPromotionList.Add(ptn);
            }
            for (int i = 0; i < removeLength; i++)
            {
                PointPromotion point = new PointPromotion(OperatorType.REMOVE)
                {
                    X = removePoints[i].X,
                    Y = removePoints[i].Y
                };
                PointPromotion ptn = ts.ApplyCoords(point, width, height);
                mPromotionList.Add(ptn);
            }
            MaskData md = this.sam.Decode(mPromotionList, this.image.ImageEmbedding, width, height);
            List<float> mask = md.mMask;
            Image<Gray, byte> resultImage = new Image<Gray, byte>(width, height);
            for (int y = 0; y < height; y++)
            {
                for (int x = 0; x < width; x++)
                {
                    int ind = y * width + x;
                    if (mask[ind] > sam.mask_threshold)
                    {
                        resultImage.Data[y, x, 0] = 255;
                    }
                }
            }
            CrazyResult result = new CrazyResult
            {
                Data = resultImage.Data,
                Width = width,
                Height = height
            };
            resultImage.Dispose();
            return result;
        }

        public CrazyResult BoxDecode(int left, int top, int right, int bottom)
        {
            int width = this.image.Width;
            int height = this.image.Height;
            List<Promotion> mPromotionList = new List<Promotion>();
            BoxPromotion promt = new BoxPromotion();
            promt.mLeftUp.X = left;
            promt.mLeftUp.Y = top;
            promt.mRightBottom.X = right;
            promt.mRightBottom.Y = bottom;
            BoxPromotion pb = ts.ApplyBox(promt, width, height);
            mPromotionList.Add(pb);
            MaskData md = this.sam.Decode(mPromotionList, this.image.ImageEmbedding, width, height);
            List<float> mask = md.mMask;
            Image<Gray, byte> resultImage = new Image<Gray, byte>(width, height);
            for (int y = 0; y < height; y++)
            {
                for (int x = 0; x < width; x++)
                {
                    int ind = y * width + x;
                    if (mask[ind] > sam.mask_threshold)
                    {
                        resultImage.Data[y, x, 0] = 255;
                    }
                }
            }
            CrazyResult result = new CrazyResult
            {
                Data = resultImage.Data,
                Width = width,
                Height = height
            };
            resultImage.Dispose();
            return result;
        }

        public bool Clear()
        {
            if (this.image != null)
            {
                this.image = null;
                this.sam.Clear();
                return true;
            }
            return false;
        }
    }
}